import os.path

import torch
import torch.utils.data as tud
from torch import nn

from model_utils import EPOCHS, device, train_and_test_model, create_all_dataloader

HIDDEN_SIZE = 3

mse_loss = nn.MSELoss().to(device)


class AutoEncoder(torch.nn.Module):
    def __init__(self, input_size: int):
        super().__init__()
        self.name = 'AutoEncoder'
        self.input_size = input_size
        self.encoder = nn.Sequential(
            nn.Linear(input_size, input_size // 2),
            nn.LeakyReLU(),
            nn.Linear(input_size // 2, input_size // 4),
            nn.LeakyReLU(),
            nn.Linear(input_size // 4, input_size // 8),
            nn.LeakyReLU(),
            nn.Linear(input_size // 8, HIDDEN_SIZE),
        )
        self.decoder = nn.Sequential(
            nn.Linear(HIDDEN_SIZE, input_size // 8),
            nn.LeakyReLU(),
            nn.Linear(input_size // 8, input_size // 4),
            nn.LeakyReLU(),
            nn.Linear(input_size // 4, input_size // 2),
            nn.LeakyReLU(),
            nn.Linear(input_size // 2, input_size),
        )

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return encoded, decoded


def get_optimizer(model):
    return torch.optim.Adam(lr=1e-4, betas=(0.9, 0.99), params=model.parameters())


def compute_auto_encoder_loss(model: nn.Module, x: torch.Tensor, **kwargs):
    x = x.view(x.shape[0], -1)
    _, decoded = model(x)
    return mse_loss(decoded, x)


def build_auto_encoder_model(input_size, **kwargs):
    model = AutoEncoder(input_size=input_size).to(device)
    optimizer = get_optimizer(model)
    return model, optimizer


def train_and_test_autoencoder_model(normal_dataloader: tud.DataLoader,
                                     attack_dataloader: tud.DataLoader,
                                     epochs: int,
                                     dataset_name: str,
                                     data_dir: str,
                                     continue_train: bool):
    return train_and_test_model(model_name='AutoEncoder',
                                compute_loss=compute_auto_encoder_loss,
                                normal_dl=normal_dataloader,
                                attack_dl=attack_dataloader,
                                epochs=epochs,
                                is_recur=False,
                                continue_train=continue_train,
                                build_model=build_auto_encoder_model,
                                dataset=dataset_name,
                                data_dir=data_dir)
